#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import logging
import subprocess

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('update_requirements')

# 原始依赖与Python 3.12兼容版本的映射
COMPATIBLE_VERSIONS = {
    'numpy==1.22.4': 'numpy==1.26.0',
    'torch==1.13.1': 'torch',  # 使用最新版本
    'transformers==4.26.1': 'transformers==4.29.2',
    'hanlp==2.1.0b36': 'hanlp>=2.1.0b36',
    'flask==2.2.3': 'flask==2.3.3',
    'werkzeug==2.2.3': 'werkzeug==2.3.7',
}

def read_requirements(file_path='requirements.txt'):
    """读取requirements.txt文件"""
    if not os.path.exists(file_path):
        logger.error(f"文件 {file_path} 不存在")
        return []
    
    with open(file_path, 'r', encoding='utf-8') as f:
        return [line.strip() for line in f if line.strip() and not line.startswith('#')]

def update_requirements(requirements):
    """更新依赖到兼容版本"""
    updated = []
    for req in requirements:
        # 检查是否需要更新此依赖
        for old_ver, new_ver in COMPATIBLE_VERSIONS.items():
            if req.lower().startswith(old_ver.split('==')[0].lower()):
                logger.info(f"更新依赖: {req} -> {new_ver}")
                req = new_ver
                break
        updated.append(req)
    return updated

def write_requirements(requirements, file_path='requirements_py312.txt'):
    """写入更新后的requirements文件"""
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write("# Updated requirements for Python 3.12 compatibility\n")
        for req in requirements:
            f.write(f"{req}\n")
    logger.info(f"已更新依赖写入到文件: {file_path}")

def check_python_version():
    """检查当前Python版本是否为3.12"""
    version = sys.version_info
    if version.major == 3 and version.minor == 12:
        logger.info("检测到Python 3.12，将更新依赖版本")
        return True
    else:
        logger.info(f"当前Python版本为{version.major}.{version.minor}，无需特别更新")
        return False

def install_updated_requirements(file_path='requirements_py312.txt'):
    """安装更新后的依赖"""
    try:
        logger.info("开始安装更新后的依赖...")
        subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-r', file_path])
        logger.info("依赖安装完成")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"安装依赖时出错: {e}")
        return False

def main():
    logger.info("开始更新requirements.txt以兼容Python 3.12")
    
    # 读取原始依赖
    requirements = read_requirements()
    if not requirements:
        return
    
    # 更新依赖版本
    updated_requirements = update_requirements(requirements)
    
    # 写入新的requirements文件
    write_requirements(updated_requirements)
    
    # 询问用户是否安装更新后的依赖
    if check_python_version():
        answer = input("是否立即安装更新后的依赖? (y/n): ").strip().lower()
        if answer == 'y':
            install_updated_requirements()
    
    logger.info("处理完成")

if __name__ == "__main__":
    main() 